# -*- coding: gb2312 -*-

from pymongo import MongoClient
from WindPy import *

import time
import datetime
import math

class StockGeneralData:
    #include all the stock data information writing and updating, e.g., stock trade status and ipo date
    def __init__(self, baseStartDateStr, baseEndDateStr = '2017-01-01',  updateDateStr = time.strftime('%Y-%m-%d'), fillType='F', periodStr = 'M', DBIP='localhost', DBPort=27017, DBName = 'stocks', periodSampleType = 'END', colIpoName = 'daily_close', colTradeStatusName = 'indices_close', colSysData = 'sys_data'):
        self.baseStartDateStr = baseStartDateStr
        self.baseEndDateStr = baseEndDateStr
        self.updateDateStr = updateDateStr
        self.fillType = fillType
        self.periodStr = periodStr
        self.DBIP = DBIP
        self.DBPort = DBPort
        self.DBName = DBName
        self.periodSampleType = periodSampleType
        self.stockIdList = []
        self.periodTransDateList = []
        self.coIpoName = colIpoName
        self.colTradeStatusName = colTradeStatusName
        self.colSysData = colSysData
        w.start()

#this is a sample mongodb operation for insert a field to a document according an index
#db.inventory.update({'qty7': { $exists: false }, 'type': 'misc'}, {$set: {'qty7': '3333'}},false,true)

    def writeBaseIpoDateToMongoDB(self, collectionName = 'ipo_date'):

        endDateStr = self.baseEndDateStr
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        colName = collectionName

        w.start()
        resStockID = w.wset('sectorconstituent', 'date=' + endDateStr + ';sectorid=a001010100000000;field=wind_code')

        windStockIdList = resStockID.Data[0]

        ipoDateDictList = []

        allStockIdStr = ''

        client = MongoClient(DBIP, DBPort)
        db = client[DBName]

        collection_ipo_date = db[collectionName]

        for i in range(0, len(windStockIdList)):
            allStockIdStr = allStockIdStr + str(windStockIdList[i]) + ','

        resIpoDate = w.wss(allStockIdStr, colName)

        stockIpoDict = {}
        for i in range(0, len(windStockIdList)) :
            aStockIdStr = str(windStockIdList[i])
            aDateStr = resIpoDate.Data[0][i].strftime('%Y-%m-%d')
            stockIpoDict[aDateStr] = aStockIdStr
        ipoDateDictList.append(stockIpoDict)
        collection_ipo_date.insert(ipoDateDictList)

        print('Finished writing base IPO date to database.')

        self.updateDataUpdateDate('ipo_date_update_date', ' ipo date ', currentDateStr = endDateStr)


    def updateIpoDateToMongoDB(self, collectionName = 'ipo_date'):


        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName

        updateDateStr = self.updateDateStr
        colName = collectionName

        w.start()
        resStockID = w.wset('sectorconstituent', 'date=' + updateDateStr + ';sectorid=a001010100000000;field=wind_code')

        windStockIdList = resStockID.Data[0]

        ipoDateDictList = []

        allStockIdStr = ''

        client = MongoClient(DBIP, DBPort)
        db = client[DBName]

        collection_ipo_date = db[collectionName]

        for i in range(0, len(windStockIdList)):
            allStockIdStr = allStockIdStr + str(windStockIdList[i]) + ','

        resIpoDate = w.wss(allStockIdStr, colName)


        for i in range(0, len(windStockIdList)):
            aStockIdStr = str(windStockIdList[i])
            aDateStr = resIpoDate.Data[0][i].strftime('%Y%m%d')
            collection_ipo_date.update({}, {'$set': {aDateStr: aStockIdStr}}, upsert=True)

        w.stop()
        print('IPO date for all stocks has been updated...')

        self.updateDataUpdateDate('ipo_date_update_date', ' IPO date for all stocks ')


    def writeBaseTradeStatusToMongoDB(self, collectionName = 'trade_status'):
        startDateStr = self.baseStartDateStr
        endDateStr = self.baseEndDateStr
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName

        w.start()
        transDateList = []

        windDateList = w.tdays(startDateStr, endDateStr).Data[0]
        for transDateObj in windDateList:
            transDateList.append(transDateObj.strftime('%Y%m%d'))

        resStockID = w.wset('sectorconstituent', 'date=' +endDateStr + ';sectorid=a001010100000000;field=wind_code')
        stockIdList = resStockID.Data[0]

        combinedStockIdStr = ''
        for stockIdStr in stockIdList :
            combinedStockIdStr = combinedStockIdStr + stockIdStr + ', '

        dateTradeStatusList = []

        #dateTradeStatusList is the stock trade status with respect to the end day of different periods, the first cell in the
        # list is the stock id list, others are the status string.
        # date string is in periodTransDateList, its length is the same as that of dateTradeStatusList (deduct 1 for the first
        # cell of the first stock id list.

        #for transDateStr in transDateList:
        windTradeStatusRes = w.wsd(combinedStockIdStr, 'trade_status', startDateStr, endDateStr, 'Fill=Previous;PriceAdj=F')
            #windTradeStatusList = windTradeStatusRes.Data[0]
            #dateTradeStatusList.append(windTradeStatusList)
        print('ok.')

        client = MongoClient(DBIP, DBPort)
        db = client[DBName]

        collection_trade_status = db[collectionName]

        tradeStatusListForDB = []
        for i in range(0, len(stockIdList)):
            tradeStatusDict = {}
            for j in range(0, len(windDateList)):
                tradeStatusDict['StockID'] = stockIdList[i]
                tradeStatusDict[transDateList[j]] = dateTradeStatusList[j][i]
            tradeStatusListForDB.append(tradeStatusDict)
            collection_trade_status.insert(tradeStatusListForDB)


    ###################################################################################################################
    ### Update Date Data as Current Date
    ###################################################################################################################
    def updateDataUpdateDate(self, strUpdatedKeyName, strUpdatedKeyAlias, currentDateStr = time.strftime('%Y-%m-%d')):
        DBIP = self.DBIP
        DBPort = self.DBPort
        DBName = self.DBName
        colSysData = self.colSysData

        client = MongoClient(DBIP, DBPort)
        db = client[DBName]
        collection_sys_data = db[colSysData]
        try:
            collection_sys_data.update({}, {'$set': {strUpdatedKeyName: currentDateStr}})
        except:
            print('Failed to update  date to MongoDB, please check.......')
            return()
        print('Successfully updated' + strUpdatedKeyAlias + 'into MongoDB.')

obj = StockGeneralData('2016-01-01')

#obj.writeBaseIpoDateToMongoDB()
#obj.updateIpoDateToMongoDB()


obj.writeBaseTradeStatusToMongoDB(obj.baseEndDateStr)
#print('Trade status data has been successfully written  to MongoDB')

#obj.writePeriodCloseToMongoDB()
#print('Period close data has been successfully written  to MongoDB')


